# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime

import pandas as pd


def get_bucket_key_1(name):
  ma_window = int(name[0:2])
  edge = int(name[4:6])
  return (ma_window, edge)


def get_bucket_key_2(name):
  ma_window = int(name[0:2])
  edge = int(name[4:6])
  wkey = 'low' if ma_window < 8 else 'mid' if ma_window < 20 else 'high'
  ekey = 'low' if edge < 8 else 'mid' if edge < 11 else 'high'
  return (wkey, ekey)


def split_into_buckets(sim_df, bucket_key_func):
  all_param_names = sorted(set(sim_df.param_name))
  param_name_buckets = {}

  for p in all_param_names:
    bkey = bucket_key_func(p)
    if bkey not in param_name_buckets:
      param_name_buckets[bkey] = []
    param_name_buckets[bkey].append(p)

  buckets = {}
  for bkey, param_names in param_name_buckets.items():
    buckets[bkey] = sim_df.loc[sim_df.param_name.isin(param_names)]

  return buckets


def filter_by_date(df, date, window=None):
  window = window if window is not None else 1
  from_dt = date - datetime.timedelta(days=window)
  return df.loc[(from_dt < df.date) & (df.date <= date)]


def get_best_n(df, num_pick, date=None, window=None):
  if date is not None:
    df = filter_by_date(df, date, window)
  df_1 = df.groupby('param_name').pnl.mean().reset_index()
  return df_1.nlargest(num_pick, 'pnl').reset_index(drop=True)


def get_worst_n(df, num_pick, date=None, window=None):
  if date is not None:
    df = filter_by_date(df, date, window)
  df_1 = df.groupby('param_name').pnl.mean().reset_index()
  return df_1.nsmallest(num_pick, 'pnl').reset_index(drop=True)


def get_mid_n(df, num_pick, date=None, window=None):
  if date is not None:
    df = filter_by_date(df, date, window)
  df_1 = df.groupby('param_name').pnl.mean().reset_index()
  df_1 = df_1.sort_by_value('pnl').reset_index(drop=True)

  num = df_1.shape[0]
  idx = max(int(num / 2. - num_pick / 2.), 0)
  return df_1.loc[idx:idx + num_pick].reset_index(drop=True)


def pick_bucket_best_n(sim_df, date, num_pick, num_pick_per_bucket, bucket_key_func):
  buckets = split_into_buckets(filter_by_date(sim_df, date), bucket_key_func)
  picked = pd.DataFrame()
  for bkey, bucket in buckets.items():
    df = get_best_n(bucket, num_pick_per_bucket)
    picked = pd.concat([picked, df], ignore_index=True)
  picked = picked.nlargest(num_pick, 'pnl').reset_index(drop=True)
  return sorted(picked['param_name'])


def pick_static(_sim_df, _date, param_names):
  return param_names


def pick_threshold(sim_df,
                   date,
                   pnl_mean_threshold=250.,
                   qty_median_threshold=80.,
                   pnl_window=10,
                   qty_window=10):
  prev_date = date - datetime.timedelta(days=1)

  # Cut by pnl
  df_1 = filter_by_date(sim_df, prev_date, pnl_window)
  pnl_df = df_1.pivot(index='date', columns='param_name', values='pnl')
  pnl_params = []
  for col in pnl_df.columns:
    if pnl_df.loc[:, col].mean() >= pnl_mean_threshold:
      pnl_params.append(str(col))

  # Cut by qty
  df_2 = filter_by_date(sim_df, prev_date, qty_window)
  qty_df = df_2.pivot(index='date', columns='param_name', values='qty_sum')
  qty_params = []
  for col in qty_df.columns:
    if qty_df.loc[:, col].median() >= qty_median_threshold:
      qty_params.append(str(col))

  return sorted(set(pnl_params) & set(qty_params))


def pick_bucket_best_n_from_historical_pnl(sim_df,
                                           date,
                                           num_pick,
                                           num_pick_per_bucket,
                                           bucket_key_func,
                                           window=1,
                                           method='best'):
  assert method in ('best', 'mid', 'worst')
  prev_date = date - datetime.timedelta(days=1)
  buckets = split_into_buckets(filter_by_date(sim_df, prev_date, window), bucket_key_func)

  picked = pd.DataFrame()
  for bkey, bucket in buckets.items():
    if method == 'worst':
      df = get_worst_n(bucket, num_pick_per_bucket)
    else:
      df = get_best_n(bucket, num_pick_per_bucket)
    picked = pd.concat([picked, df], ignore_index=True)

  if picked.empty:
    return []

  if method == 'best':
    picked = picked.nlargest(num_pick, 'pnl').reset_index(drop=True)
  elif method == 'mid':
    df = picked.sort_values('pnl').reset_index(drop=True)
    idx = max(int(df.shape[0] / 2. - num_pick / 2.), 0)
    picked = df.loc[idx:idx + num_pick - 1].reset_index(drop=True)
  else:
    picked = picked.nsmallest(num_pick, 'pnl').reset_index(drop=True)

  return sorted(picked['param_name'])


def pick_by_historical_pnl(sim_df, date, num_pick, window=1, method='best'):
  assert method in ('best', 'mid', 'worst')
  prev_date = date - datetime.timedelta(days=1)
  df = filter_by_date(sim_df, prev_date, window)

  df_1 = df.groupby('param_name')['pnl'].mean().reset_index()
  df_1['group'] = [n[:8] for n in df_1.param_name]
  df_1 = df_1.set_index('param_name')

  grp = df_1.groupby('group')
  df_2 = pd.DataFrame([{'param_name': df.pnl.idxmax(), 'pnl': df.pnl.max()} for _, df in grp])
  if df_2.empty:
    return []

  if method == 'best':
    picked = df_2.nlargest(num_pick, 'pnl').reset_index(drop=True)
  elif method == 'mid':
    df = df_2.sort_values('pnl').reset_index(drop=True)
    idx = max(int(df.shape[0] / 2. - num_pick / 2.), 0)
    picked = df.loc[idx:idx + num_pick - 1].reset_index(drop=True)
  else:
    picked = df_2.nsmallest(num_pick, 'pnl').reset_index(drop=True)
  return sorted(picked['param_name'])


def pick_bucket_worst_by_reversion(sim_df,
                                   date,
                                   num_pick,
                                   num_pick_per_bucket,
                                   bucket_key_func,
                                   loss_threshold=-2000):
  prev_date = date - datetime.timedelta(days=1)
  buckets = split_into_buckets(filter_by_date(sim_df, prev_date), bucket_key_func)

  picked = pd.DataFrame()
  for bkey, bucket in buckets.items():
    df = bucket.loc[bucket.pnl < loss_threshold]
    picked = pd.concat([picked, df], ignore_index=True)

  picked = picked.nsmallest(num_pick, 'pnl').reset_index(drop=True)
  return sorted(picked['param_name'])


def pick_by_min_corr(sim_df, date, params, corr_mat=None):
  if corr_mat is None:
    pnl_df = sim_df.pivot(index='date', columns='param_name', values='pnl')
    corr_mat = pnl_df.corr()

  picked_params = set()
  for param in params:
    picked = corr_mat[param].nsmallest(len(params)).reset_index()
    for param in picked.param_name:
      if param not in params and param not in picked_params:
        picked_params.add(param)
        break
  return sorted(picked_params)


def pick_bucket_mid(sim_df, date, num_pick=14, num_per_bucket=3):
  assert num_pick == 14

  # Universe
  universe = pick_threshold(sim_df,
                            date,
                            pnl_mean_threshold=250,
                            qty_median_threshold=65,
                            pnl_window=40,
                            qty_window=10)
  refined_df = sim_df.loc[sim_df.param_name.isin(universe)]

  names = pick_bucket_best_n_from_historical_pnl(refined_df,
                                                 date,
                                                 100,
                                                 1,
                                                 get_bucket_key_1,
                                                 window=60,
                                                 method='best')
  refined_df = refined_df.loc[refined_df.param_name.isin(names)]

  return pick_bucket_best_n_from_historical_pnl(refined_df,
                                                date,
                                                num_pick,
                                                num_per_bucket,
                                                get_bucket_key_2,
                                                window=1,
                                                method='mid')


def pick_bucket_worst(sim_df, date, num_pick=14, num_per_bucket=3):
  assert num_pick == 14

  # Universe
  universe = pick_threshold(sim_df,
                            date,
                            pnl_mean_threshold=250,
                            qty_median_threshold=65,
                            pnl_window=40,
                            qty_window=10)
  refined_df = sim_df.loc[sim_df.param_name.isin(universe)]

  names = pick_bucket_best_n_from_historical_pnl(refined_df,
                                                 date,
                                                 100,
                                                 1,
                                                 get_bucket_key_1,
                                                 window=60,
                                                 method='best')
  refined_df = refined_df.loc[refined_df.param_name.isin(names)]

  return pick_bucket_best_n_from_historical_pnl(refined_df,
                                                date,
                                                num_pick,
                                                num_per_bucket,
                                                get_bucket_key_2,
                                                window=1,
                                                method='worst')


def pick_hybrid(sim_df, date, num_pick=14):
  assert num_pick == 14
  prev_date = date - datetime.timedelta(days=1)

  names = pick_bucket_best_n_from_historical_pnl(sim_df,
                                                 date,
                                                 9,
                                                 2,
                                                 get_bucket_key_2,
                                                 window=1,
                                                 method='best')
  names += pick_by_historical_pnl(sim_df, date, 9, window=8, method='worst')
  names = set(names)

  df = sim_df.loc[sim_df.param_name.isin(names)]
  names = pick_bucket_best_n_from_historical_pnl(df,
                                                 date,
                                                 100,
                                                 1,
                                                 get_bucket_key_1,
                                                 window=20,
                                                 method='best')
  df = df.loc[df.param_name.isin(names)]

  return sorted(get_best_n(df, num_pick, prev_date, window=1).param_name)


def pick_hybrid_2(sim_df, date, num_pick=14):
  universe_hiqty = pick_threshold(sim_df,
                                  date,
                                  pnl_mean_threshold=100,
                                  qty_median_threshold=200,
                                  pnl_window=30,
                                  qty_window=10)
  refined_df_hiqty = sim_df.loc[sim_df.param_name.isin(universe_hiqty)]
  names_hiqty = pick_bucket_best_n_from_historical_pnl(refined_df_hiqty,
                                                       date,
                                                       100,
                                                       1,
                                                       get_bucket_key_1,
                                                       window=5,
                                                       method='best')
  refined_df_hiqty = refined_df_hiqty.loc[refined_df_hiqty.param_name.isin(names_hiqty)]
  names_hiqty = pick_by_historical_pnl(refined_df_hiqty,
                                       date,
                                       num_pick // 2,
                                       window=1,
                                       method='worst')

  universe_loqty = pick_threshold(sim_df,
                                  date,
                                  pnl_mean_threshold=500,
                                  qty_median_threshold=20,
                                  pnl_window=60,
                                  qty_window=60)
  refined_df_loqty = sim_df.loc[sim_df.param_name.isin(universe_loqty)]
  names_loqty = pick_bucket_best_n_from_historical_pnl(refined_df_loqty,
                                                       date,
                                                       100,
                                                       1,
                                                       get_bucket_key_1,
                                                       window=5,
                                                       method='best')
  refined_df_loqty = refined_df_loqty.loc[refined_df_loqty.param_name.isin(names_loqty)]
  names_loqty = pick_by_historical_pnl(refined_df_loqty,
                                       date,
                                       num_pick // 2,
                                       window=2,
                                       method='best')

  return sorted(set(names_hiqty + names_loqty))


def pick_hybrid_3(sim_df, date, num_pick=14):
  assert num_pick == 14

  low_ma_df = sim_df.loc[sim_df.ma_window <= 2]
  low_ma = pick_bucket_best_n_from_historical_pnl(low_ma_df,
                                                  date,
                                                  100,
                                                  1,
                                                  get_bucket_key_1,
                                                  window=15,
                                                  method='best')
  refined_df_lowma = low_ma_df.loc[low_ma_df.param_name.isin(low_ma)]
  names = pick_by_historical_pnl(refined_df_lowma, date, 5, window=1, method='best')

  universe_loqty = pick_threshold(sim_df,
                                  date,
                                  pnl_mean_threshold=500,
                                  qty_median_threshold=15,
                                  pnl_window=60,
                                  qty_window=60)
  refined_df_loqty = sim_df.loc[sim_df.param_name.isin(universe_loqty)]
  names_loqty = pick_bucket_best_n_from_historical_pnl(refined_df_loqty,
                                                       date,
                                                       100,
                                                       1,
                                                       get_bucket_key_1,
                                                       window=5,
                                                       method='best')
  refined_df_loqty = refined_df_loqty.loc[refined_df_loqty.param_name.isin(names_loqty)]

  for i in range(4, 11):
    names_loqty = pick_by_historical_pnl(refined_df_loqty, date, i, window=2, method='best')
    if len(set(names + names_loqty)) >= 10:
      break
  names = sorted(set(names + names_loqty))

  universe_hiqty = pick_threshold(sim_df,
                                  date,
                                  pnl_mean_threshold=100,
                                  qty_median_threshold=150,
                                  pnl_window=30,
                                  qty_window=10)
  refined_df_hiqty = sim_df.loc[sim_df.param_name.isin(universe_hiqty)]
  names_hiqty = pick_bucket_best_n_from_historical_pnl(refined_df_hiqty,
                                                       date,
                                                       100,
                                                       1,
                                                       get_bucket_key_1,
                                                       window=5,
                                                       method='best')
  refined_df_hiqty = refined_df_hiqty.loc[refined_df_hiqty.param_name.isin(names_hiqty)]

  for i in range(4, 15):
    names_hiqty = pick_by_historical_pnl(refined_df_hiqty, date, i, window=1, method='worst')
    if len(set(names + names_hiqty)) >= 14:
      break
  names = sorted(set(names + names_hiqty))

  return names
